#pragma once

#include "stdafx.h"
#include "FileImage.h"
#include "MemImage.h"
#include "Base.h"

#define MAX_FUNCTION_NAME 0x50

typedef struct _SSDT_ENTRIES
{
    ULONG Index;
    ULONG CurrentAddress;
    ULONG OrgAddress;
    TCHAR Name[MAX_FUNCTION_NAME];
    TCHAR OrgModuleName[MAX_PATH];
    TCHAR CurrentModuleName[MAX_PATH];
}SSDT_ENTRIES, *PSSDT_ENTRIES;

typedef struct _KSERVICE_TABLE_DESCRIPTOR {
    PULONG_PTR Base;
    PULONG Count;
    ULONG Limit;
    PUCHAR Number;
} KSERVICE_TABLE_DESCRIPTOR, *PKSERVICE_TABLE_DESCRIPTOR;

class CSSDT:public CEnumPECallback, CBase
{
public:
    CSSDT(void);
    CSSDT(HANDLE hDriver);
    ~CSSDT(void);

    virtual BOOL EnumRelocationDataCallBack(
            IN CPEParse* clsPeFile,
            IN PIMAGE_BASE_RELOCATION pRelocBlock,
            IN ULONG RelocAddress
        );
    virtual BOOL EnumExportFunctionCallBack(
            IN CPEParse* clsPeFile,
            IN LPCSTR lpszFuncName,
            IN ULONG Address,
            IN ULONG FunctionId,
            IN LPCSTR lpszForward
        );

    virtual BOOL EnumImportFunctionCallBack()
    {
        return TRUE;
    }

    BOOL
    QuerySSDTHookInfo(
        VOID
    );

    BOOL
    QueryShadowSSDTHookInfo(
        VOID
    );

private:

    ULONG
    GetOrgSSDTEntryCounts(
        VOID
    );

    ULONG
    GetOrgShadowSSDTEntryCounts(
        VOID
    );

    BOOL
    InitGetOrgSSDTEntries(
        VOID
    );

    BOOL
    InitGetOrgShadowSSDTEntries(
        VOID
    );

    BOOL 
    UpdateSSDTFunctionName(
        VOID
    );

    BOOL 
    UpdateShadowSSDTFunctionName(
        VOID
    );

    BOOL
    UpdateOrgSSDTEntries(
        VOID
    );

    BOOL
    UpdateOrgShadowSSDTEntries(
        VOID
    );

    BOOL 
    UpdateCurrentSSDTEntries(
        VOID
    );

    BOOL 
    UpdateCurrentShadowSSDTEntries(
        VOID
    );

    ULONG 
    GetCurrentSSDTEntryConnts(
        VOID
    );

    ULONG 
    GetCurrentShadowSSDTEntryConnts(
        VOID
    );

private:
    ULONG m_KeServiceDescTableRva;
    ULONG m_KiServiceTableRva;
    ULONG m_KiServiceLimit;
    ULONG m_Win32ServiceTableRva;
    ULONG m_Win32ServiceLimit;
    ULONG m_SSDTEntryCounts;
    PSSDT_ENTRIES m_SSDTHookInfo;
    ULONG m_ShadowSSDTEntryCounts;
    PSSDT_ENTRIES m_ShadowSSDTHookInfo;
};

